package com.toymate.ai.llm.model.dify;

import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.alibaba.fastjson.JSON;
import com.toymate.common.utils.PushUtils;
import com.toymate.modules.iot.entity.IotMessage;
import com.toymate.modules.iot.entity.IotRole;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;

import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * DifySSEEventSourceListener
 *
 * @author weirdor (179946422@qq.com)
 * @version v1.0.0
 */
@Slf4j
public class DifySSEEventSourceListener extends EventSourceListener {

    private static final Pattern SENTENCE_END_PATTERN = Pattern.compile("[。！？；?.!;]");
    private static PushUtils pushUtils;

    static {
        DifySSEEventSourceListener.pushUtils = SpringUtil.getBean(PushUtils.class);
    }

    private StringBuilder answer = new StringBuilder();
    private StringBuilder totalAnswer = new StringBuilder();
    private int count;
    private AtomicInteger index = new AtomicInteger(0);
    private String productId;
    private String deviceNum;
    private String conversationId;
    private String question;
    private IotRole role;

    public DifySSEEventSourceListener(String productId, String deviceNum, String question, IotRole role) {
        this.productId = productId;
        this.deviceNum = deviceNum;
        this.question = question;
        this.role = role;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void onOpen(EventSource eventSource, Response response) {
        log.info("======== dify 对话开始 ========");
    }

    /**
     * {@inheritDoc}
     */
    @SneakyThrows
    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
        //处理数据并合成语音
        DifyCompletionResponse response = JSON.parseObject(data, DifyCompletionResponse.class);
        if ("message_end".equals(response.getEvent())) {
            if (StrUtil.isNotEmpty(answer.toString())) {
                log.info("结束数据：{} 序号:{}", answer, index);
                pushUtils.pushMessage(productId, deviceNum, answer.toString(), index.get(), role);
                answer.setLength(0);
                this.conversationId = response.getConversation_id();
                log.info("会话ID：{}",this.conversationId);
            }
            log.info("======== 对话结束 ========");
        }
        if ("message".equals(response.getEvent())) {
            String content = response.getAnswer().trim();
            if (StrUtil.isNotEmpty(content)) {
                answer.append(content);
                totalAnswer.append(content);
                Matcher matcher = SENTENCE_END_PATTERN.matcher(answer.toString());
                int lastMatchPosition = -1;
                while (matcher.find()) {
                    lastMatchPosition = matcher.end();
                }
                if (lastMatchPosition != -1) {
                    if (lastMatchPosition >= count) {
                        String message = answer.substring(0, lastMatchPosition);
                        log.info("数据：{} 序号:{}", message, index.get());
                        pushUtils.pushMessage(productId, deviceNum, message, index.get(), role);
                        answer.delete(0, lastMatchPosition);
                        count = Math.min(count + 30, 80);
                        index.incrementAndGet();
                    }
                }
            }

        }
    }

    @Override
    public void onClosed(@NotNull EventSource eventSource) {
        IotMessage iotMessage = new IotMessage();
        iotMessage.setDeviceId(this.deviceNum);
        iotMessage.setQuestion(this.question);
        iotMessage.setAnswer(this.totalAnswer.toString());
        iotMessage.setConversationId(conversationId);
        iotMessage.insert();
    }

    @SneakyThrows
    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        if (Objects.isNull(response)) {
            return;
        }
        ResponseBody body = response.body();
        if (Objects.nonNull(body)) {
            log.error("Dify  sse连接异常data：{}，异常：{}", body.string(), t);
        } else {
            log.error("Dify  sse连接异常data：{}，异常：{}", response, t);
        }
        eventSource.cancel();
    }

}
